package com.example.emos.wx.config.xss;

import cn.hutool.core.util.StrUtil;
import cn.hutool.http.HtmlUtil;
import cn.hutool.json.JSONUtil;

import javax.servlet.ReadListener;
import javax.servlet.ServletInputStream;
import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletRequestWrapper;
import java.io.*;
import java.nio.charset.Charset;
import java.util.LinkedHashMap;
import java.util.Map;

/**
 * XssHttpServletRequestWrapper 类是对 HttpServletRequestWrapper 的扩展，用于处理请求参数和请求体中的 XSS 攻击。
 * 继承父类 HttpServletRequestWrapper，重写相关方法进行参数和请求体的处理。
 */
public class XssHttpServletRequestWrapper extends HttpServletRequestWrapper {

    /**
     * 构造方法，接收原始 HttpServletRequest 对象，并调用父类构造器进行初始化。
     *
     * @param request 原始 HttpServletRequest 对象
     */
    public XssHttpServletRequestWrapper(HttpServletRequest request) {
        super(request);
    }

    /**
     * 重写 getParameter 方法，对请求参数进行 XSS 过滤。
     *
     * @param name 参数名
     * @return 过滤后的参数值
     */
    @Override
    public String getParameter(String name) {
        // 获取原始的参数值
        String value = super.getParameter(name);
        // 判断参数值是否为空，如果不为空则进行 XSS 过滤
        if (!StrUtil.hasEmpty(value)) {
            value = HtmlUtil.filter(value);
        }
        return value;
    }

    /**
     * 重写 getParameterValues 方法，对请求参数数组进行 XSS 过滤。
     *
     * @param name 参数名
     * @return 过滤后的参数值数组
     */
    @Override
    public String[] getParameterValues(String name) {
        // 获取原始的参数值数组
        String[] values = super.getParameterValues(name);
        // 如果数组不为空，遍历数组进行 XSS 过滤
        if (values != null) {
            for (int i = 0; i < values.length; i++) {
                String value = values[i];
                // 判断数组元素是否为空，如果不为空则进行 XSS 过滤
                if (!StrUtil.hasEmpty(value)) {
                    value = HtmlUtil.filter(value);
                }
                // 将过滤后的元素重新放回数组中
                values[i] = value;
            }
        }
        return values;
    }

    /**
     * 重写 getParameterMap 方法，对请求参数 Map 进行 XSS 过滤。
     *
     * @return 过滤后的请求参数 Map
     */
    @Override
    public Map<String, String[]> getParameterMap() {
        // 获取原始的请求参数 Map
        Map<String, String[]> parameters = super.getParameterMap();
        // 创建一个新的 LinkedHashMap 用于存放过滤后的参数
        LinkedHashMap<String, String[]> map = new LinkedHashMap();
        // 遍历原始参数 Map 的所有键
        if (parameters != null) {
            for (String key : parameters.keySet()) {
                // 获取键对应的值数组
                String[] values = parameters.get(key);
                // 遍历值数组进行 XSS 过滤
                for (int i = 0; i < values.length; i++) {
                    String value = values[i];
                    // 判断数组元素是否为空，如果不为空则进行 XSS 过滤
                    if (!StrUtil.hasEmpty(value)) {
                        value = HtmlUtil.filter(value);
                    }
                    // 将过滤后的元素重新放回数组中
                    values[i] = value;
                }
                // 将过滤后的键值对放入新的 Map 中
                map.put(key, values);
            }
        }
        return map;
    }

    /**
     * 重写 getHeader 方法，对请求头进行 XSS 过滤。
     *
     * @param name 请求头名称
     * @return 过滤后的请求头值
     */
    @Override
    public String getHeader(String name) {
        // 获取原始的请求头值
        String value = super.getHeader(name);
        // 判断请求头值是否为空，如果不为空则进行 XSS 过滤
        if (!StrUtil.hasEmpty(value)) {
            value = HtmlUtil.filter(value);
        }
        return value;
    }

    /**
     * 重写 getInputStream 方法，对请求体进行 XSS 过滤。
     *
     * @return 过滤后的请求体的 ServletInputStream
     * @throws IOException IO 异常
     */
    @Override
    public ServletInputStream getInputStream() throws IOException {
        // 获取原始的输入流
        InputStream in = super.getInputStream();

        // 将原始输入流包装成 InputStreamReader，指定字符集为UTF-8
        InputStreamReader reader = new InputStreamReader(in, Charset.forName("UTF-8"));

        // 将 InputStreamReader 包装成 BufferedReader，提高读取效率
        BufferedReader buffer = new BufferedReader(reader);

        // 使用 StringBuilder 存储读取的请求体内容
        StringBuilder body = new StringBuilder();

        // 逐行读取输入流中的数据，并拼接成字符串
        String line = buffer.readLine();
        while (line != null) {
            body.append(line);
            line = buffer.readLine();
        }

        // 关闭流
        buffer.close();
        reader.close();
        in.close();

        // 将请求体转换为 Map 对象
        Map<String, Object> map = JSONUtil.parseObj(body.toString());
        Map<String, Object> result = new LinkedHashMap<>();

        // 遍历 Map，对字符串值进行 XSS 过滤
        for (String key : map.keySet()) {
            Object val = map.get(key);
            // 判断val是否是字符串
            if (val instanceof String) {
                if (!StrUtil.hasEmpty(val.toString())) {
                    // 对字符串进行 XSS 过滤
                    result.put(key, HtmlUtil.filter(val.toString()));
                }
            } else {
                // 如果值不是字符串，直接添加到结果 Map 中
                result.put(key, val);
            }
        }

        // 将过滤后的 Map 转换为 JSON 字符串
        String json = JSONUtil.toJsonStr(result);

        // 将 JSON 字符串转换为 ByteArrayInputStream
        ByteArrayInputStream bain = new ByteArrayInputStream(json.getBytes());

        // 返回包装后的 ServletInputStream
        return new ServletInputStream() {
            @Override
            public int read() throws IOException {
                return bain.read();
            }

            @Override
            public boolean isFinished() {
                return false;
            }

            @Override
            public boolean isReady() {
                return false;
            }

            @Override
            public void setReadListener(ReadListener readListener) {
                // 这里未实现异步读取监听器
            }
        };
    }
}

